"""
Jaeger分布式追踪服务
提供应用内分布式追踪功能
"""

import os
from typing import Optional
from opentelemetry import trace
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from fastapi import FastAPI
from sqlalchemy.orm import Session
from src.utils.logging import get_logger

logger = get_logger(__name__)


class TracingService:
    """分布式追踪服务类"""
    
    def __init__(self):
        self.tracer = None
        self.tracer_provider = None
        self.enabled = os.getenv("ENABLE_TRACING", "false").lower() == "true"
        logger.info(f"分布式追踪服务初始化，启用状态: {self.enabled}")
    
    def setup_tracing(self, service_name: str = "langChainGraph"):
        """
        设置Jaeger分布式追踪
        
        Args:
            service_name: 服务名称
        """
        if not self.enabled:
            logger.info("分布式追踪已禁用")
            return
        
        try:
            # 设置服务资源
            resource = Resource(attributes={
                SERVICE_NAME: service_name
            })
            
            # 创建TracerProvider
            self.tracer_provider = TracerProvider(resource=resource)
            trace.set_tracer_provider(self.tracer_provider)
            
            # 配置Jaeger导出器
            jaeger_exporter = JaegerExporter(
                agent_host_name=os.getenv("JAEGER_AGENT_HOST", "localhost"),
                agent_port=int(os.getenv("JAEGER_AGENT_PORT", "6831")),
            )
            
            # 添加批处理Span处理器
            span_processor = BatchSpanProcessor(jaeger_exporter)
            self.tracer_provider.add_span_processor(span_processor)
            
            # 获取tracer实例
            self.tracer = trace.get_tracer(__name__)
            
            logger.info("Jaeger分布式追踪设置完成")
        except Exception as e:
            logger.error(f"设置分布式追踪失败: {e}")
    
    def instrument_fastapi(self, app: FastAPI):
        """
        为FastAPI应用添加追踪仪器
        
        Args:
            app: FastAPI应用实例
        """
        if not self.enabled or not self.tracer_provider:
            return
            
        try:
            FastAPIInstrumentor.instrument_app(app)
            logger.info("FastAPI应用追踪仪器设置完成")
        except Exception as e:
            logger.error(f"为FastAPI应用设置追踪仪器失败: {e}")
    
    def instrument_requests(self):
        """为requests库添加追踪仪器"""
        if not self.enabled or not self.tracer_provider:
            return
            
        try:
            RequestsInstrumentor().instrument()
            logger.info("requests库追踪仪器设置完成")
        except Exception as e:
            logger.error(f"为requests库设置追踪仪器失败: {e}")
    
    def instrument_sqlalchemy(self, engine):
        """
        为SQLAlchemy引擎添加追踪仪器
        
        Args:
            engine: SQLAlchemy引擎实例
        """
        if not self.enabled or not self.tracer_provider:
            return
            
        try:
            SQLAlchemyInstrumentor().instrument(engine=engine)
            logger.info("SQLAlchemy引擎追踪仪器设置完成")
        except Exception as e:
            logger.error(f"为SQLAlchemy引擎设置追踪仪器失败: {e}")
    
    def create_span(self, name: str, attributes: Optional[dict] = None):
        """
        创建追踪span
        
        Args:
            name: Span名称
            attributes: Span属性
            
        Returns:
            Span对象
        """
        if not self.enabled or not self.tracer:
            return None
            
        span = self.tracer.start_span(name)
        if attributes:
            for key, value in attributes.items():
                span.set_attribute(key, value)
        return span
    
    def get_tracer(self):
        """
        获取tracer实例
        
        Returns:
            tracer实例
        """
        return self.tracer if self.enabled else None


# 全局追踪服务实例
tracing_service = TracingService()